from PIL import Image
from transformers import SamModel, SamProcessor
class SlimSAM():
    def __init__(self,experiments_dir,logger,device='npu'):
        self.logger = logger
        self.device = device
        self.logger.info("Loading SlimSAM model...")
        self.model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-50", cache_dir=experiments_dir).to(device)
        self.processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-50", cache_dir=experiments_dir)
        self.logger.info(f"SlimSAM model loaded successfully. Device: {self.device}")
    def preprocess_image(self, image_path,bbox):
        # 检查bbox需要时1维数组
        assert len(bbox)==4, "bbox should be a 1-d array with 4 elements"
        
        raw_image=Image.open(image_path)
        inputs = self.processor(raw_image, input_boxes=[bbox], return_tensors="pt").to(self.device)
        outputs = self.model(**inputs)
        masks = self.processor.image_processor.post_process_masks(
            outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
        return masks[0]
    def visualize_mask(self, image_path, mask,save_path):
        raw_image = Image.open(image_path)
        mask = Image.fromarray(mask)
        raw_image.paste(mask, (0, 0), mask)
        if save_path:
            raw_image.save(save_path)
            self.logger.info(f"Mask saved to {save_path}")
            